import tushare as ts
import pandas as pd
import time
import os
import multiprocessing
import datetime


os.chdir(r'**')
token = '*8b334e2d5d482292aa8b516c7f*59dd6d23c0266dac3dc407dfb63ea1*'
pro = ts.pro_api(token)

day1 = datetime.date(2021, 1, 1).strftime('%Y%m%d')
if datetime.date.today().isoweekday() > 5:
    today = datetime.date.today() - datetime.timedelta(days=1)
    today = today.strftime('%Y%m%d')
else:
    today = datetime.date.today().strftime('%Y%m%d')


def tedliu(ts_code_p):
    print(ts_code_p)
    df = pro.daily(ts_code=ts_code_p, start_date=day1, end_date=today, fields='ts_code,trade_date,pct_chg')
    df.replace("2021", "", regex=True, inplace=True)
    df = df.T
    df.columns = df.iloc[1]
    df.drop("trade_date", inplace=True)
    df.rename(index={"pct_chg": ts_code_p}, inplace=True)
    df.drop("ts_code", inplace=True)
    return df


if __name__ == "__main__":
    print("...Starting...")
    now = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    print(now)
    now1 = time.time()
    data = pro.query('stock_basic', exchange='', list_status='L', fields='ts_code,name,industry')
    pool = multiprocessing.Pool(processes=8)
    multiple_results = [pool.apply_async(tedliu, args=(i,)) for i in data['ts_code']]
    lst = [res.get() for res in multiple_results]
    pool.close()
    pool.join()
    result = pd.concat(lst)
    result.reset_index(inplace=True)
    result.rename(columns={'index': 'ts_code'}, inplace=True)
    result = pd.merge(result, data, on=["ts_code"])

    ts_name = result['name']
    result.drop(labels=['name'], axis=1, inplace=True)
    result.insert(0, 'name', ts_name)

    industry = result['industry']
    result.drop(labels=['industry'], axis=1, inplace=True)
    result.insert(1, 'industry', industry)

    pe = pro.query('daily_basic', ts_code='', trade_date=today, fields='ts_code, pe_ttm')
    result = pd.merge(result, pe, on=["ts_code"])
    pe_name = result['pe_ttm']
    result.drop(labels=['pe_ttm'], axis=1, inplace=True)
    result.insert(2, 'pe_ttm', pe_name)

    result.to_excel('daily.xlsx')

    print("------------------------------------")
    print("...Finished...")
    now = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    print(now)
    now2 = time.time()
    now3 = (now2 - now1)/60
    print("...用时{} Mins...".format(round(now3, 2)))
